# -*- coding:utf-8 -*-
# Author:凌逆战 | Never
# Date: 2023/1/1
"""
绘制 对比不同的窗函数和和对应的频率响应
"""
import numpy as np
from numpy.fft import rfft
import matplotlib.pyplot as plt
from window_function import Rectangle_windows, Voibis_windows, Orka_forward_window, Orka_backward_window

window_len = 128
hop_size = 32
NFFT = 256


def frequency_response(window, window_len=window_len, NFFT=NFFT):
    A = rfft(window, NFFT) / (window_len / 2)  # (513,)
    mag = np.abs(A)
    freq = np.linspace(0, 0.5, len(A))
    # 忽略警告
    with np.errstate(divide='ignore', invalid='ignore'):
        response = 20 * np.log10(mag)
    response = np.clip(response, -150, 150)
    return freq, response


# Rectangle_windows = Rectangle_windows(window_len)
hanning_window = np.hanning(M=window_len)
hamming_window = np.hamming(M=window_len)
Voibis_windows = Voibis_windows(window_len)
blackman_window = np.blackman(M=window_len)
bartlett_window = np.bartlett(M=window_len)
kaiser_window = np.kaiser(M=window_len, beta=14)
N1 = hop_size
N2 = window_len - hop_size
Orka_forward_window = Orka_forward_window(N1=N1, N2=N2, hop_size=hop_size, window_len=window_len)

plt.figure()
# plt.plot(Rectangle_windows, label="Rectangle")
plt.plot(hanning_window, label="hanning")
# plt.plot(hamming_window, label="hamming")
plt.plot(Voibis_windows, label="Voibis")
plt.plot(blackman_window, label="blackman")
# plt.plot(bartlett_window, label="bartlett")
plt.plot(kaiser_window, label="kaiser")
plt.plot(Orka_forward_window, label="Orka_forward")

plt.legend()
plt.tight_layout()
plt.show()

# freq, Rectangle_FreqResp = frequency_response(Rectangle_windows, window_len)
freq, hanning_FreqResp = frequency_response(hanning_window, window_len)
freq, hamming_FreqResp = frequency_response(hamming_window, window_len)
freq, Voibis_FreqResp = frequency_response(Voibis_windows, window_len)
freq, blackman_FreqResp = frequency_response(blackman_window, window_len)
freq, bartlett_FreqResp = frequency_response(bartlett_window, window_len)
freq, kaiser_FreqRespw = frequency_response(kaiser_window, window_len)
freq, Orka_forward_window = frequency_response(Orka_forward_window, window_len)

plt.figure()
plt.title("Frequency response")
# plt.plot(freq, Rectangle_FreqResp, label="Rectangle")
plt.plot(freq, hanning_FreqResp, label="hanning")
# plt.plot(freq, hamming_FreqResp, label="hamming")
plt.plot(freq, Voibis_FreqResp, label="Voibis")
plt.plot(freq, blackman_FreqResp, label="blackman")
# plt.plot(freq, bartlett_FreqResp, label="bartlett")
plt.plot(freq, kaiser_FreqRespw, label="kaiser")
plt.plot(freq, Orka_forward_window, label="Orka_forward")
plt.ylabel("Magnitude [dB]")
plt.xlabel("Normalized frequency [cycles per sample]")

plt.legend()
plt.tight_layout()
plt.show()
